# -*- coding: utf-8 -*-
"""
    Project name：code_potentialflow
# -------------------------------
    File name：mobius.py
    Created on：2025/4/12 16:30
    Author：(input)
    Description:
"""
import sympy as sp
import numpy as np


class MobiusTransformSolver:
    def __init__(self, z_vals, w_vals):
        """
        初始化：接受 4 个 z 点及其对应的 w 值（符号也可含 beta）
        z_vals: list of 4 complex or real numbers
        w_vals: list of 4 expressions or values (can contain symbol beta)
        """
        self.z_vals = z_vals
        self.w_vals = w_vals
        self.z = sp.Symbol('z', complex=True)  # 支持复数输入
        self.a, self.b, self.d, self.beta = sp.symbols('a b d beta', real=True)
        self.f_expr = (self.a * self.z + self.b) / (self.z + self.d)
        self.solution = None

    def solve(self):
        """
        求解 Möbius 变换表达式，自动筛选满足 0 < beta < 1 的解
        在求解成功时打印参数 a, b, d, beta 和表达式
        返回：满足条件的解字典（包含 a, b, d, beta）
        """
        equations = [
            sp.Eq(self.f_expr.subs(self.z, z), w)
            for z, w in zip(self.z_vals, self.w_vals)
        ]
        solutions = sp.solve(equations, (self.a, self.b, self.d, self.beta), dict=True)

        for sol in solutions:
            beta_val = sol[self.beta]
            if sp.simplify(sp.And(beta_val > 0, beta_val < 1)) == True:
                self.solution = sol
                print("✅ 求解成功：")
                print(f"a = {sp.N(sol[self.a], 10)}")
                print(f"b = {sp.N(sol[self.b], 10)}")
                print(f"d = {sp.N(sol[self.d], 10)}")
                print(f"beta = {sp.N(sol[self.beta], 10)}")

                print("\n Möbius 变换表达式：")
                f_expr = self.get_expression()
                sp.pprint(sp.Eq(sp.Function('f')(self.z), f_expr))
                ipt_a = sp.N(sol[self.a], 10)
                ipt_b = sp.N(sol[self.b], 10)
                ipt_d = sp.N(sol[self.d], 10)
                ipt_beta = sp.N(sol[self.beta], 10)

                return sol, ipt_a, ipt_b, ipt_d, ipt_beta

        print("未找到满足 beta ∈ (0,1) 的解")
        return None

    def get_expression(self):
        """
        返回求解后的 Möbius 表达式
        """
        if self.solution is None:
            raise ValueError("尚未求解，或未找到有效解")
        return self.f_expr.subs(self.solution)

    def verify(self):
        """
        验证初始 4 个点是否满足映射表达式
        """
        if self.solution is None:
            raise ValueError("尚未求解，无法验证")
        f_final = self.get_expression()
        print("\n验证初始映射点：")
        for z, w_expected in zip(self.z_vals, self.w_vals):
            w_calc = sp.N(f_final.subs(self.z, z), 10)
            if isinstance(w_expected, (sp.Expr, sp.Symbol)):
                w_target = sp.N(w_expected.subs(self.solution), 10)
            else:
                w_target = sp.N(w_expected, 10)
            print(f"f({z}) = {w_calc},  目标: {w_target}")

    def map_point(self, z_value):
        """
        对任意 z 值使用求得的 Möbius 变换表达式进行映射
        """
        if self.solution is None:
            raise ValueError("尚未求解，无法进行映射")
        f_final = self.get_expression()
        return sp.N(f_final.subs(self.z, z_value), 10)

    def transform_coordinates(self, points):
        """
        批量变换坐标点，可选择放缩到指定范围
        参数：
            points: 列表，包含字典 {'label': str, 'u': complex}
        返回：
            列表，包含原始点、放缩坐标和变换后的坐标
        """
        if self.solution is None:
            raise ValueError("尚未求解，无法进行映射")

        results = []
        for point in points:
            label = point['label']
            z = point['z']

            z_mapped = self.map_point(z)

            results.append({
                'label': label,
                'z': z,
                'z_mapped': z_mapped
            })

        return results


if __name__ == "__main__":
    # 初始化点对
    z_vals = [1, 1.6, 17813, -17813]
    beta = sp.Symbol('beta', real=True)
    w_vals = [-1 / beta, -1, 1, 1 / beta]

    # 实例化并求解
    solver = MobiusTransformSolver(z_vals, w_vals)
    sol, a, b, d, Beta = solver.solve()# 返回Mobius相关的几个参数

    if sol:
        solver.verify()

        # 示例：计算任意点的映射
        z_test = 0
        w_test = solver.map_point(z_test)
        print(f"\n任意点 z = {z_test} 映射结果为：w = {w_test}")

        # 测试 Mobius 变换
        print("\n测试 Mobius 变换：")
        point_data = [
            {"label": "A", "z": -1},
            {"label": "B", "z": -8696},
            {"label": "C", "z": -17813},
            {"label": "D", "z": 8696},
            {"label": "E", "z": 1.6},
            {"label": "F", "z": 1},
            {"label": "J", "z": 17813}
        ]

        # 批量变换坐标
        results = solver.transform_coordinates(point_data)

        # 打印结果
        print("\n角点、控制点的Mobius变换：")
        for entry in results:
            print(f"点 {entry['label']}：")
            print(f"  原坐标 z     = {entry['z']}")
            print(f"  映射后 z'     = {entry['z_mapped']}")
            print("-" * 40)
    else:
        print("未找到满足 beta ∈ (0,1) 的解")